# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/6 下午4:16
@Author yzpang
@Function: 
"""
import numpy as np
from modelserver.configs.base_config import get_logger
from modelserver.configs.server_config import MULTI_CLASS, MULTI_LABEL
from sklearn.metrics import f1_score, classification_report


class MetricReport:
    """
    F1 score for text classification task.
    """

    def __init__(self, name="MetricReport", average="micro", classification_type=MULTI_CLASS):
        super(MetricReport, self).__init__()
        self.average = average
        self._name = name
        self.classification_type = classification_type
        self.reset()

    def reset(self):
        """
        Resets all of the metric state.
        """
        self.y_prob = None
        self.y_true = None

    def f1_score(self, y_prob):
        """
        Compute micro f1 score and macro f1 score
        """
        # 多分类
        if self.classification_type == MULTI_CLASS:
            self.y_pred = np.argmax(y_prob, axis=-1)
        # 层次分类
        elif self.classification_type == MULTI_LABEL:
            self.y_pred = y_prob > 0.5
        micro_f1_score = f1_score(y_pred=self.y_pred, y_true=self.y_true, average="micro")
        macro_f1_score = f1_score(y_pred=self.y_pred, y_true=self.y_true, average="macro")
        return micro_f1_score, macro_f1_score

    def update(self, probs, labels):
        """
        Update the probability and label
        """
        if self.y_prob is not None:
            self.y_prob = np.append(self.y_prob, probs.numpy(), axis=0)
        else:
            self.y_prob = probs.numpy()
        if self.y_true is not None:
            self.y_true = np.append(self.y_true, labels.numpy(), axis=0)
        else:
            self.y_true = labels.numpy()

    def accumulate(self):
        """
        Returns micro f1 score and macro f1 score
        """
        micro_f1_score, macro_f1_score = self.f1_score(y_prob=self.y_prob)
        return micro_f1_score, macro_f1_score

    def report(self):
        """
        Returns classification report
        """
        # 多分类
        if self.classification_type == MULTI_CLASS:
            self.y_pred = np.argmax(self.y_prob, axis=-1)
        # 层次分类
        elif self.classification_type == MULTI_LABEL:
            self.y_pred = self.y_prob > 0.5
        get_logger().info("classification report:\n" + classification_report(self.y_true, self.y_pred, digits=4))

    def name(self):
        """
        Returns metric name
        """
        return self._name